﻿#include <learn/alg/alg_coeff_matrix.h>

using namespace xi;
using namespace xi::alg;

Eigen::SparseMatrix<double> Alg_CoeffMatrix::curve_least_square(const std::vector<double> &params,
                                                                const std::vector<double> &knots, std::size_t p)
{
    auto m = params.size() - 1;
    auto n = knots.size() - p - 2;

    std::vector<Eigen::Triplet<double>> tripletsB0, tripletsB1, tripletsB2;
    for (std::size_t i = 0; i < m + 1; i++)
    {
        double t = params[i];
        auto span = find_span(knots.data(), knots.size(), p, t);
        auto basis = basis_funcs(knots.data(), p, span, t);

        for (std::size_t j = 0; j < p + 1; j++)
            tripletsB0.emplace_back(i, span - p + j, basis[j]);
    }

    Eigen::SparseMatrix<double> B0(m + 1, n + 1);
    B0.setFromTriplets(tripletsB0.begin(), tripletsB0.end());
    return B0;
}

std::array<Eigen::SparseMatrix<double>, 2> Alg_CoeffMatrix::curve_stretch_bending(const std::vector<double> &knots,
                                                                                  std::size_t p)
{
    // 2 个基函数相乘，要确保精确，至少需要 p 阶
    static Integrator_Gauss integrator(p);

    auto n = knots.size() - p - 2;

    // 计算 B1 * B1^T, B2 * B2^T 的积分
    auto &weights = integrator.weights;
    auto &nodes = integrator.nodes;
    std::vector<Eigen::Triplet<double>> tripletsB1, tripletsB2;
    for (std::size_t i = p; i < n + 1; i++)
    {
        double mid = (knots[i + 1] + knots[i]) / 2;
        double ratio = (knots[i + 1] - knots[i]) / 2;

        // 防御重节点，此时区间为零，不用积分
        if (std::abs(ratio) > std::numeric_limits<double>::epsilon())
        {
            for (std::size_t j = 0; j < nodes.size(); j++)
            {
                auto derivs = deriv_basis_funcs(knots.data(), p, i, 2, mid + nodes[j] * ratio);
                for (std::size_t ii = 0; ii < p + 1; ii++)
                {
                    for (std::size_t jj = 0; jj < p + 1; jj++)
                    {
                        tripletsB1.emplace_back(i - p + ii, i - p + jj, derivs[1][ii] * derivs[1][jj] * weights[j]);
                        tripletsB2.emplace_back(i - p + ii, i - p + jj, derivs[2][ii] * derivs[2][jj] * weights[j]);
                    }
                }
            }
        }
    }

    Eigen::SparseMatrix<double> B1(n + 1, n + 1), B2(n + 1, n + 1);
    B1.setFromTriplets(tripletsB1.begin(), tripletsB1.end());
    B2.setFromTriplets(tripletsB2.begin(), tripletsB2.end());

    return {B1, B2};
}

Eigen::SparseMatrix<double> Alg_CoeffMatrix::surface_least_square(const std::vector<Vec2> &params,
                                                                  const std::vector<double> &uKnots,
                                                                  const std::vector<double> &vKnots, std::size_t p,
                                                                  std::size_t q)
{
    auto m = params.size() - 1;
    auto n1 = uKnots.size() - p - 2;
    auto n2 = vKnots.size() - q - 2;

    // 最小二乘矩阵 B: (m + 1) x ((n1 + 1) x (n2 + 1))
    std::vector<Eigen::Triplet<double>> tripletsB0;
    for (std::size_t i = 0; i < m + 1; i++)
    {
        Vec2 P = params[i];
        auto uSpan = find_span(uKnots.data(), uKnots.size(), p, P.x());
        auto vSpan = find_span(vKnots.data(), vKnots.size(), q, P.y());

        auto ubasis = basis_funcs(uKnots.data(), p, uSpan, P.x());
        auto vbasis = basis_funcs(vKnots.data(), q, vSpan, P.y());

        for (std::size_t j = 0; j < p + 1; j++)
            for (std::size_t k = 0; k < q + 1; k++)
                tripletsB0.emplace_back(i, (uSpan - p + j) * (n2 + 1) + (vSpan - q + k), ubasis[j] * vbasis[k]);
    }

    auto N = (n1 + 1) * (n2 + 1);
    Eigen::SparseMatrix<double> B0(m + 1, N);
    B0.setFromTriplets(tripletsB0.begin(), tripletsB0.end());
    return B0;
}

std::array<Eigen::SparseMatrix<double>, 2> Alg_CoeffMatrix::surface_stretch_bending(const std::vector<double> &uKnots,
                                                                                    const std::vector<double> &vKnots,
                                                                                    std::size_t p, std::size_t q)
{
    // 4 个基函数相乘，要确保精确，至少需要 p + q 阶
    static Integrator_Gauss integrator(p + q);

    auto n1 = uKnots.size() - p - 2;
    auto n2 = vKnots.size() - q - 2;

    // 计算基函数的向量
    Eigen::MatrixXd Bu0(n1 + 1, n1 + 1), Bu1(n1 + 1, n1 + 1), Bu2(n1 + 1, n1 + 1);
    Eigen::MatrixXd Bv0(n2 + 1, n2 + 1), Bv1(n2 + 1, n2 + 1), Bv2(n2 + 1, n2 + 1);
    Bu0.setZero();
    Bu1.setZero();
    Bu2.setZero();
    Bv0.setZero();
    Bv1.setZero();
    Bv2.setZero();

    auto &weights = integrator.weights;
    auto &nodes = integrator.nodes;
    for (std::size_t i = p; i < n1 + 1; i++)
    {
        double mid = (uKnots[i + 1] + uKnots[i]) / 2;
        double ratio = (uKnots[i + 1] - uKnots[i]) / 2;

        // 防御重节点，此时区间为零，不用积分
        if (std::abs(ratio) > std::numeric_limits<double>::epsilon())
        {
            for (std::size_t j = 0; j < nodes.size(); j++)
            {
                auto derivs = deriv_basis_funcs(uKnots.data(), p, i, 2, mid + nodes[j] * ratio);
                for (std::size_t ii = 0; ii < p + 1; ii++)
                {
                    for (std::size_t jj = 0; jj < p + 1; jj++)
                    {
                        // 注意权重要缩放一下，考虑到积分区间长度不同
                        Bu0(i - p + ii, i - p + jj) += derivs[0][ii] * derivs[0][jj] * weights[j] * ratio;
                        Bu1(i - p + ii, i - p + jj) += derivs[1][ii] * derivs[1][jj] * weights[j] * ratio;
                        Bu2(i - p + ii, i - p + jj) += derivs[2][ii] * derivs[2][jj] * weights[j] * ratio;
                    }
                }
            }
        }
    }

    for (std::size_t i = q; i < n2 + 1; i++)
    {
        double mid = (vKnots[i + 1] + vKnots[i]) / 2;
        double ratio = (vKnots[i + 1] - vKnots[i]) / 2;

        // 防御重节点，此时区间为零，不用积分
        if (std::abs(ratio) > std::numeric_limits<double>::epsilon())
        {
            for (std::size_t j = 0; j < nodes.size(); j++)
            {
                auto derivs = deriv_basis_funcs(vKnots.data(), q, i, 2, mid + nodes[j] * ratio);
                for (std::size_t ii = 0; ii < q + 1; ii++)
                {
                    for (std::size_t jj = 0; jj < q + 1; jj++)
                    {
                        // 注意权重要缩放一下，考虑到积分区间长度不同
                        Bv0(i - q + ii, i - q + jj) += derivs[0][ii] * derivs[0][jj] * weights[j] * ratio;
                        Bv1(i - q + ii, i - q + jj) += derivs[1][ii] * derivs[1][jj] * weights[j] * ratio;
                        Bv2(i - q + ii, i - q + jj) += derivs[2][ii] * derivs[2][jj] * weights[j] * ratio;
                    }
                }
            }
        }
    }

    std::vector<Eigen::Triplet<double>> tripletsB1, tripletsB2;
    for (std::size_t i = 0; i < n1 + 1; i++)
    {
        for (std::size_t j = 0; j < n2 + 1; j++)
        {
            auto row = i * (n2 + 1) + j;
            for (std::size_t ii = 0; ii < n1 + 1; ii++)
            {
                for (std::size_t jj = 0; jj < n2 + 1; jj++)
                {
                    auto col = ii * (n2 + 1) + jj;

                    double Eu = Bu1(i, ii) * Bv0(j, jj);
                    double Ev = Bu0(i, ii) * Bv1(j, jj);

                    double Euu = Bu2(i, ii) * Bv0(j, jj);
                    double Euv = Bu1(i, ii) * Bv1(j, jj);
                    double Evv = Bu0(i, ii) * Bv2(j, jj);

                    tripletsB1.emplace_back(row, col, Eu + Ev);
                    tripletsB2.emplace_back(row, col, Euu + Euv * 2 + Evv);
                }
            }
        }
    }

    auto N = (n1 + 1) * (n2 + 1);
    Eigen::SparseMatrix<double> B1(N, N), B2(N, N);
    B1.setFromTriplets(tripletsB1.begin(), tripletsB1.end());
    B2.setFromTriplets(tripletsB2.begin(), tripletsB2.end());

    return {B1, B2};
}
